package main

import (
	"encoding/csv"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

type DialogData struct {
	Input  string `json:"input"`
	Target string `json:"target"`
}

// 解析target字段，提取意图名称
func extractIntent(target string) string {
	var targetData []map[string]interface{}
	if err := json.Unmarshal([]byte(target), &targetData); err != nil {
		return "解析失败"
	}

	if len(targetData) > 0 {
		for key := range targetData[0] {
			return key
		}
	}
	return "未知意图"
}

// 从input中提取客户问题
func extractQuestion(input string) string {
	start := strings.Index(input, "客户：")
	if start == -1 {
		return "未找到问题"
	}

	questionPart := input[start+len("客户："):]
	end := strings.Index(questionPart, "\n")
	if end == -1 {
		end = len(questionPart)
	}

	return strings.TrimSpace(questionPart[:end])
}

// 从文件读取JSON数据
func readJSONFromFile(filename string) ([]DialogData, error) {
	file, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	var data []DialogData
	if err := json.Unmarshal(file, &data); err != nil {
		return nil, err
	}

	return data, nil
}

func main() {
	// 方法1: 使用硬编码数据
	//jsonData := `[
	//	{
	//		"input": "基于客户与坐席的对话历史，综合判断客户最后一句话所属的意图和槽位。\n###对话历史开始###\n客户：有个自称淘宝客服的打电话说我订单异常，要我退款，是真的假的？\n###对话历史结束###\n客户最后一句话所属的意图与槽位是：",
	//		"target": "[{\"咨询反诈\": {}}]"
	//	},
	//	{
	//		"input": "基于客户与坐席的对话历史，综合判断客户最后一句话所属的意图和槽位。\n###对话历史开始###\n客户：警察说我涉嫌洗钱，要我配合调查，是真的吗？\n###对话历史结束###\n客户最后一句话所属的意图与槽位是：",
	//		"target": "[{\"咨询反诈\": {}}]"
	//	}
	//]`
	//
	//var data []DialogData
	//if err := json.Unmarshal([]byte(jsonData), &data); err != nil {
	//	fmt.Printf("解析JSON失败: %v\n", err)
	//	return
	//}

	// 方法2: 如果要从文件读取，取消注释下面的代码

	data, err := readJSONFromFile("C:\\Users\\nanjiang6\\Downloads\\train_2.105.json")
	if err != nil {
		fmt.Printf("从文件读取JSON失败: %v\n", err)
		return
	}

	// 创建CSV文件
	file, err := os.Create("dialog_data.csv")
	if err != nil {
		fmt.Printf("创建CSV文件失败: %v\n", err)
		return
	}
	defer file.Close()

	// 创建CSV writer，支持中文
	writer := csv.NewWriter(file)
	writer.UseCRLF = true // Windows系统换行符
	defer writer.Flush()

	// 写入UTF-8 BOM头，确保Excel正确显示中文
	bom := []byte{0xEF, 0xBB, 0xBF}
	if _, err := file.Write(bom); err != nil {
		fmt.Printf("写入BOM头失败: %v\n", err)
		return
	}

	// 写入表头
	header := []string{"意图名称", "客户问题"}
	if err := writer.Write(header); err != nil {
		fmt.Printf("写入表头失败: %v\n", err)
		return
	}

	// 写入数据行
	for _, item := range data {
		intent := extractIntent(item.Target)
		question := extractQuestion(item.Input)

		record := []string{intent, question}
		if err := writer.Write(record); err != nil {
			fmt.Printf("写入数据失败: %v\n", err)
			return
		}
	}

	fmt.Println("数据已成功写入 dialog_data.csv 文件")

	// 显示统计信息
	fmt.Printf("\n成功处理 %d 条记录\n", len(data))

	// 预览前几条记录
	fmt.Println("\n前几条记录预览：")
	fmt.Printf("%-20s | %s\n", "意图名称", "客户问题")
	fmt.Println(strings.Repeat("-", 80))
	for i, item := range data {
		if i >= 5 { // 只显示前5条
			break
		}
		intent := extractIntent(item.Target)
		question := extractQuestion(item.Input)
		fmt.Printf("%-20s | %s\n", intent, question)
	}
}
